import turtle
import numpy as np
import math
from common import pprint, exec_time


def on_close(func):
    """
    窗口关闭时，关闭程序
    """
    def wrapper(*args, **kwargs):
        try:
            pprint('开始绘制')
            func(*args, **kwargs)
            pprint('绘制完成')
            turtle.done()
        except Exception:
            pprint('绘制终止')
    return wrapper
    

def move_turtle(x, y):
    """
    抬起海龟移动
    x: x轴坐标
    y: y轴坐标
    """
    turtle.up()
    turtle.goto(x, y)
    turtle.down()


@on_close
@exec_time
def draw_image(data, width=1920, height=1080):
    turtle.title('按行绘制图片')
    turtle.setworldcoordinates(0, height+5, width+5, 0)  # 设置坐标
    turtle.ht()

    for row in range(height):
        # start_time = time.time()
        # 绘制一行
        pprint(f'图像绘制中{format(row/height*100, ".2f")}%。。。。', end='\r')
        move_turtle(0, row)

        for col in range(width):
            # 绘制一个像素点
            turtle.color(tuple(data[row][col][:3]))
            turtle.fd(1)
        turtle.update()

@on_close
@exec_time
def draw_col_image(data, width=1920, height=1080):
    turtle.title('按列绘制图片')
    turtle.setworldcoordinates(0, height+5, width+5, 0)  # 设置坐标
    turtle.ht()
    turtle.lt(90)

    for col in range(width):
        # 绘制一列
        pprint(f'图像绘制中{format(col/width*100, ".2f")}%。。。。', end='\r')
        move_turtle(col, 0)

        for row in range(height):
            # 绘制一个像素点
            turtle.color(tuple(data[row][col]))
            turtle.fd(1)
        turtle.update()


@on_close
@exec_time
def draw_chroma(data, width=1920, height=1080):
    turtle.title('绘制回形图案')
    # turtle.setworldcoordinates(-(width+10)//2, -height//2, (width+10)//2, height//2)
    turtle.ht()

    dist = 1
    next_dist = 1
    step = 1
    count = 1
    total = data.shape[0]
    for item in data:
        turtle.color(tuple(item))
        turtle.fd(1)
        dist -= 1
        if dist == 0:
            pprint(f'图像绘制中{format(count/total*100, ".2f")}%。。。。', end='\r')
            dist = next_dist
            count += dist
            next_dist += step
            step = 1 - step
            turtle.rt(90)
            turtle.update()


@on_close
@exec_time
def draw_swirl(data, width=1920, height=1080):
    turtle.title('绘制漩涡图案')
    turtle.setworldcoordinates(-(width+10)//2, -height//2, (width+10)//2, height//2)
    turtle.ht()
    
    dist = 2
    index = 2
    count = 1
    total = data.shape[0]
    for item in data:
        turtle.color(tuple(item))
        turtle.fd(1)
        dist -= 1
        if dist <= 0:
            pprint(f'图像绘制中{format(count/total*100, ".2f")}%。。。。', end='\r')
            turtle.lt(91)
            dist = index * 1.5
            count += dist
            index += 1
            turtle.update()


def draw(data, mode='row', width=1920, height=1080):
    """
    data: 图像数据
    mode: 绘制图像模式 ['row', 'col', 'chroma', 'swirl']
    """
    turtle.tracer(False)
    screen = turtle.Screen()
    screen.colormode(255)
    screen.bgcolor("black")

    if mode == 'row':
        turtle.setup(width+5, height+5)
        draw_image(data, width, height)
    elif mode == 'col':
        turtle.setup(width+5, height+5)
        draw_col_image(data, width, height)
    elif mode == 'chroma':
        # width = height = int(math.sqrt(width*height))
        turtle.setup(width+5, height+5)
        draw_chroma(data, width, height)
    elif mode == 'swirl':
        turtle.setup(width, height)
        draw_swirl(data, width, height)
